
import cv2
import numpy as np
from .myostu import lmt_ostu, lmt_ostu2



def watershed(image):

    # 1、将图像转为灰度图像
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 2、阈值分割,将图像分为黑白两部分
    ret, thresh = lmt_ostu2(gray, 95, True)

    # 3、对图像进行“开运算”,先腐蚀再膨胀
    kernel = np.ones((3, 3), np.uint8)
    opening = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kernel, iterations=2)

    # 4、对“开运算”的结果进行膨胀,得到大部分都是背景的区域
    sure_bg = cv2.dilate(opening, kernel, iterations=3)

    # 5、通过distanceTransform获取前景区域
    dist_transform = cv2.distanceTransform(opening, cv2.DIST_L2, 5)
    ret, sure_fg = cv2.threshold(dist_transform, 0.1*dist_transform.max(), 255, 0)

    # 6、sure_bg与sure_fg相减,得到既有前景又有背景的重合区域
    sure_fg = np.uint8(sure_fg)
    unknow = cv2.subtract(sure_bg, sure_fg)

    # 7、连通区域处理
    # 对连通区域进行标号
    ret, markers = cv2.connectedComponents(sure_fg, connectivity=8)

    # OpenCV 分水岭算法对物体做的标注必须都 大于1, 背景标号为0
    markers = markers + 1

    markers[unknow == 255] = 0

    # 8、分水岭算法
    markers = cv2.watershed(image, markers)  # 分水岭算法后, 所有轮廓的像素点被标注为 -1

    mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
    mask[markers == 1] = 255

    return mask


